# -*- coding: utf-8 -*-
# @Time    : 2020/6/17 下午10:08
# @Author  : caotian
# @FileName: optimizationtrain.py
# @Software: PyCharm
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D,Pool2D,Linear
import numpy as np
import json
import gzip
import random
import os
import sys
from PIL import Image
curpath=os.path.abspath(os.curdir)
sys.path.append(curpath)
import optimizationdata as od
import optimizationmodel as om
train_loader=od.load_data('train')
use_gpu=False
place=fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()
with fluid.dygraph.guard(place):
    model=om.MNIST()
    model.train()
    optimizer=fluid.optimizer.SGDOptimizer(learning_rate=0.01,parameter_list=model.parameters())
    epoch_num=5
    for epoch_id in range(epoch_num):
        for batch_id,data in enumerate(train_loader()):
            image_data,label_data=data
            image=fluid.dygraph.to_variable(image_data)
            label=fluid.dygraph.to_variable(label_data)
            # 前向计算的过程，同时拿到模型输出值和分类准确率
            if batch_id == 0 and epoch_id == 0:
                # 打印模型参数和每层输出的尺寸
                predict, acc = model(image, label, check_shape=True, check_content=False)
            elif batch_id == 401:
                # 打印模型参数和每层输出的值
                predict, acc = model(image, label, check_shape=False, check_content=True)
            else:
                predict, acc = model(image, label)

            loss=fluid.layers.cross_entropy(predict,label)
            avg_loss=fluid.layers.mean(loss)
            if batch_id % 200==0:
                print("epoch:{},batch:{},loss :{}".format(epoch_id,batch_id,avg_loss.numpy()))
            avg_loss.backward()
            optimizer.minimize(avg_loss)
            model.clear_gradients()
    fluid.save_dygraph(model.state_dict(),'mnist-model')
